import hashlib
import hmac
import base64
from pathlib import Path
from typing import Union
from cryptography.fernet import Fernet
from ..core.exceptions import SecurityError

class Security:
    def __init__(self, secret_key: str = None):
        if secret_key:
            self.fernet = Fernet(secret_key.encode())
        else:
            self.fernet = Fernet(Fernet.generate_key())

    def encrypt_data(self, data: bytes) -> bytes:
        """加密数据"""
        try:
            return self.fernet.encrypt(data)
        except Exception as e:
            raise SecurityError(f"加密失败: {str(e)}")

    def decrypt_data(self, encrypted_data: bytes) -> bytes:
        """解密数据"""
        try:
            return self.fernet.decrypt(encrypted_data)
        except Exception as e:
            raise SecurityError(f"解密失败: {str(e)}")

    @staticmethod
    def verify_signature(data: bytes, signature: str, key: str) -> bool:
        """验证签名"""
        try:
            expected_signature = hmac.new(
                key.encode(),
                data,
                hashlib.sha256
            ).hexdigest()
            return hmac.compare_digest(signature, expected_signature)
        except Exception as e:
            raise SecurityError(f"签名验证失败: {str(e)}")

    @staticmethod
    def calculate_checksum(file_path: Union[str, Path]) -> str:
        """计算文件校验和"""
        try:
            sha256_hash = hashlib.sha256()
            with open(file_path, "rb") as f:
                for byte_block in iter(lambda: f.read(4096), b""):
                    sha256_hash.update(byte_block)
            return sha256_hash.hexdigest()
        except Exception as e:
            raise SecurityError(f"计算校验和失败: {str(e)}")

    @staticmethod
    def verify_checksum(file_path: Union[str, Path], expected_checksum: str) -> bool:
        """验证文件校验和"""
        actual_checksum = Security.calculate_checksum(file_path)
        return hmac.compare_digest(actual_checksum, expected_checksum)